from keras.layers import Input, Conv2D, DepthwiseConv2D, Activation, Reshape, Flatten, \
                         Dense, concatenate, BatchNormalization, MaxPool2D, Lambda
from keras.activations import relu
from keras.models import Model
from keras.utils import plot_model

import tensorflow as tf
import numpy as np

def Relu6(x, max_value=6):
    return relu(x, max_value=max_value)

def Conv(x, filter, filter_shape = (3,3), stride=2, has_relu = True):
    _x = Conv2D(filter, filter_shape, strides=stride, padding="same")(x)
    _x = BatchNormalization()(_x)
    if has_relu: _x = Activation("relu")(_x)
    return _x

def dsConv(x, filter_shape = (3, 3), stride=2):
    _x = DepthwiseConv2D(filter_shape, strides=stride, padding="same")(x)
    _x = BatchNormalization()(_x)
    _x = Activation("relu")(_x)
    return _x

def backbone(input, detect_branch=[]):
    x = input

    x = dsConv(x, filter_shape = (10,4), stride=2)
    x = Conv(x, 64, (1, 1), 1)
    
    x = dsConv(x, stride=1)  
    x = Conv(x, 64, (1, 1), 1)
    
    x = dsConv(x, stride=1)
    x = Conv(x, 64, (1, 1), 1)
    
    x = dsConv(x, stride=1)
    x = Conv(x, 64, (1,1), 1)
    
    x = dsConv(x, stride=1)
    x = Conv(x, 64, (1,1), 1)
    
    return Model(input, x)

def base_model1(shape, class_num = 10):
    back_bone_model = backbone(Input(shape=shape, name="input_1"))
    from keras.layers import GlobalAveragePooling2D, Flatten, Dense
    x = GlobalAveragePooling2D()(back_bone_model.output)
    x = Dense(class_num)(x)
    x = Activation('softmax', name="softmax_1")(x)
    
    return Model(back_bone_model.input, outputs = x)